Skip to content

Support containers in to/from_numpy - #178

Merged
inducer merged 8 commits into
inducer:mainfrom
alexfikl:actx-to-numpy-containers
Jul 1, 2022
Merged

Support containers in to/from_numpy#178
inducer merged 8 commits into
inducer:mainfrom
alexfikl:actx-to-numpy-containers

Conversation

@alexfikl

@alexfikl alexfikl commented Jun 26, 2022 •
edited
Loading

Copy link
Copy Markdown
Collaborator

This is a bit bigger than it's supposed to be because it moves some things around to avoid a bit of duplicated code, but the main changes are:

  • adds support for containers to actx.to_numpy and actx.from_numpy (deprecating the functions in traversal.py)
  • adds support for containers to actx.empty_like (and actx.zeros_like) to avoid confusion with actx.np.empty_like (and actx.np.zeros_like).
  • adds some helpers for TaggableCLArray.
  • Fix jax.np.all and jax.np.any: it was just forwarding to jnp.all which doesn't work right on containers.

With this, all the functions in ArrayContext that should work with containers do.

TODO

  • Just forward actx.empty_like in actx.np.empty_like (same for zeros_like)
  • Add some tests? Not sure any more are needed since the functions are tested in other places with containers.

@alexfikl
alexfikl force-pushed the actx-to-numpy-containers branch 6 times, most recently from 1cfa1d4 to 1ea2620 Compare June 27, 2022 08:04
@alexfikl alexfikl mentioned this pull request Jun 27, 2022
Open
@alexfikl
alexfikl marked this pull request as ready for review June 27, 2022 08:10
@inducer
inducer force-pushed the actx-to-numpy-containers branch from 1ea2620 to f4ab671 Compare June 28, 2022 02:34
inducer
inducer reviewed Jun 28, 2022
View reviewed changes

@inducer inducer left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for working on this! A few comments below.

Comment thread arraycontext/impl/pyopencl/__init__.py Outdated
Comment thread arraycontext/impl/pyopencl/__init__.py Outdated
from arraycontext.impl.pyopencl.fake_numpy import PyOpenCLFakeNumpyNamespace
return PyOpenCLFakeNumpyNamespace(self)

def _map_container(self, func, array, allowed_types=None, strict=False):

@inducer inducer Jun 28, 2022

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Type annotation?

@alexfikl alexfikl Jun 28, 2022

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in https://github.com/inducer/arraycontext/compare/f4ab671ce9eb96e9dd2397d9a77d5e7ac7c9be36..971df793aed6b5b107be09588c0791a2f664c2b9. I think it's correct, but probably could use a serious look.

Comment thread arraycontext/impl/jax/__init__.py Outdated
Comment thread arraycontext/impl/pyopencl/__init__.py
Comment thread arraycontext/impl/jax/__init__.py
Comment thread arraycontext/impl/pytato/__init__.py Outdated
Comment thread arraycontext/impl/pytato/__init__.py Outdated
Comment thread arraycontext/impl/pytato/__init__.py Outdated
Comment thread arraycontext/fake_numpy.py
Comment thread arraycontext/impl/jax/fake_numpy.py Outdated
@alexfikl
alexfikl force-pushed the actx-to-numpy-containers branch 5 times, most recently from 85a96ec to 2dcdbfe Compare June 29, 2022 07:25
@alexfikl

alexfikl commented Jun 29, 2022

Copy link
Copy Markdown
Collaborator Author

@inducer This should be ready for another go. I thought we could leave the deprecation parts for a follow-up PR, since this is already confusing and moves a lot of stuff around.

@alexfikl
alexfikl force-pushed the actx-to-numpy-containers branch 3 times, most recently from 66df690 to 049c1e5 Compare June 29, 2022 17:36
@inducer

inducer commented Jun 29, 2022

Copy link
Copy Markdown
Owner

Uh-oh, sorry. This has picked up some conflicts. 😬

@alexfikl
alexfikl force-pushed the actx-to-numpy-containers branch 2 times, most recently from c7c102a to c1f557c Compare June 30, 2022 06:31
@alexfikl
alexfikl force-pushed the actx-to-numpy-containers branch from c1f557c to bb4d8fc Compare June 30, 2022 06:42
@alexfikl

alexfikl commented Jun 30, 2022

Copy link
Copy Markdown
Collaborator Author

Uh-oh, sorry. This has picked up some conflicts.

Should be fixed now! Hopefully didn't miss anything.

@alexfikl
alexfikl requested a review from inducer June 30, 2022 17:15
inducer
inducer reviewed Jun 30, 2022
View reviewed changes

@inducer inducer left a comment •
edited
Loading

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Just a few final things, then this is good to go from my end.

Comment thread arraycontext/context.py Outdated
Comment thread arraycontext/context.py Outdated
Comment thread arraycontext/context.py Outdated
Comment thread arraycontext/impl/pytato/__init__.py
@alexfikl
alexfikl force-pushed the actx-to-numpy-containers branch from bb4d8fc to 8c9952a Compare July 1, 2022 06:28
@alexfikl
alexfikl requested a review from inducer July 1, 2022 06:32
@inducer
inducer enabled auto-merge (rebase) July 1, 2022 19:58
@inducer

inducer commented Jul 1, 2022

Copy link
Copy Markdown
Owner

LGTM, thanks!

@inducer
inducer merged commit fd4d4d8 into inducer:main Jul 1, 2022
@inducer

inducer commented Jul 1, 2022

Copy link
Copy Markdown
Owner

I think this might have broken pytential: https://github.com/inducer/loopy/runs/7156901597?check_suite_focus=true 😬

@alexfikl alexfikl mentioned this pull request Jul 2, 2022
Merged
@alexfikl
alexfikl deleted the actx-to-numpy-containers branch July 2, 2022 04:35
@alexfikl

alexfikl commented Jul 2, 2022

Copy link
Copy Markdown
Collaborator Author

I think this might have broken pytential: inducer/loopy/runs/7156901597?check_suite_focus=true grimacing

Oops! Sorry about that.. should be fixed by inducer/pytential#166!

@inducer

inducer commented Jul 2, 2022

Copy link
Copy Markdown
Owner

No need to apologize! Thanks for pushing a fix quickly.

This was referenced Jul 2, 2022
Merged
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@inducer inducer Awaiting requested review from inducer

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@alexfikl @inducer